import time
import matplotlib
import matplotlib.animation as animation
import matplotlib.pyplot as plt
import nengo
import numpy as np
from PIL import Image
from collections import deque
from multiprocessing import Pool, Process
from xml.etree import ElementTree as et
from alive_progress import alive_bar
# Access to enums and MuJoCo library functions.
from dm_control.mujoco.wrapper.mjbindings import enums, mjlib
from dm_control import composer, mjcf, mujoco
from dm_control.utils import inverse_kinematics as ik
from IPython.display import HTML
%env MUJOCO_GL=egl
%matplotlib widget
env: MUJOCO_GL=egl
# Rendering parameters
dpi = 100
framerate = 60 # (Hz)
width, height = 720, 480
sensor_shape = (10, 10)
# IK solver parameters
_MAX_STEPS = 100
_TOL = 1e-12
_TIME_STEP = 2e-3
def display_video(frames, framerate=framerate, figsize=None):
height, width, _ = frames[0].shape
orig_backend = matplotlib.get_backend()
matplotlib.use('Agg') # Switch to headless 'Agg' to inhibit figure rendering.
if figsize is None:
figsize = (width / dpi, height / dpi)
fig, ax = plt.subplots(1, 1, figsize=figsize, dpi=dpi)
matplotlib.use(orig_backend) # Switch back to the original backend.
ax.set_axis_off()
ax.set_aspect('equal')
ax.set_position([0, 0, 1, 1])
im = ax.imshow(frames[0])
def update(frame):
im.set_data(frame)
return [im]
interval = 1000/framerate
anim = animation.FuncAnimation(fig=fig, func=update, frames=frames,
interval=interval, blit=True, repeat=False)
return HTML(anim.to_html5_video())
def J_zero(physics, site_name):
phys = mujoco.Physics.from_model(physics.model)
jac_pos = np.zeros((3, phys.model.nv))
jac_rot = np.zeros((3, phys.model.nv))
mjlib.mj_jacSite(
phys.model.ptr,
phys.data.ptr,
jac_pos,
jac_rot,
phys.model.name2id(site_name, 'site'))
return np.vstack((jac_pos, jac_rot))
# Define 3D arrays of balls connected with sliders
_M = (15, 15)
radius = 5e-4
mass = 1e-12/np.prod(_M)
dx = (3e-2-2*radius)/(_M[0] - 1)
offset = dx*(_M[0] - 1)/2.0
geom_type = 'sphere'
robot_xml = 'models/panda_nohand.xml'
tree = et.parse(robot_xml)
root = tree.getroot()
for body in root.findall('.//body'):
if body.get('name') == 'sensor_pad':
sensor_pad = body
# If balls and tendons are already defined, remove them first
if sensor_pad.find('body') is not None:
for body in sensor_pad.findall('body'):
sensor_pad.remove(body)
tendon = root.find('tendon')
if tendon is not None:
for spatial in tendon.findall('spatial'):
tendon.remove(spatial)
def create_linked_ball(sid, pos, euler, springdamper=(1e-2, 1)):
# Create base ball
elem = et.Element('body')
elem.set('name', 'taxel' + str(sid))
elem.set('pos', '{:.6f} {:.6f} {:.6f}'.format(pos[0], pos[1], pos[2]))
elem.set('euler', euler)
# Inertia
inertia = et.Element('inertial')
inertia.set('mass', str(mass))
inertia.set('pos', '0 0 0')
inertia.set('fullinertia', '1e-6 1e-6 1e-6 0 0 0')
# Joint
joint = et.Element('joint')
joint.set('type', 'slide')
joint.set('axis', '0 0 1')
joint.set('springdamper', '{} {}'.format(springdamper[0], springdamper[1]))
joint.set('frictionloss', '1e-2')
# Geom
geom = et.Element('geom')
geom.set('name', 'ball' + str(sid))
geom.set('type', geom_type)
geom.set('size', str(radius))
geom.set('class', 'collision')
# Site
site = et.Element('site')
site.set('name', 'tendon_site' + str(sid))
site.set('type', 'sphere')
site.set('size', '1e-4')
site.set('group', '1')
# Add to dynamic tree
elem.append(inertia)
elem.append(joint)
elem.append(geom)
elem.append(site)
return elem
def create_tendon(tid, site_names):
spatial = et.Element('spatial')
spatial.set('name', 'tendom'+str(tid))
spatial.set('width', '2e-4')
spatial.set('limited', 'true')
spatial.set('range', '0 {:.6f}'.format(4*dx))
spatial.set('frictionloss', '1e-2')
spatial.set('stiffness', '0.1')
spatial.set('damping', '0.5')
for sn in site_names:
site = et.Element('site')
site.set('site', sn)
spatial.append(site)
return spatial
# Define balls with sliders and touch sensors
sid = 0
for i in range(_M[0]):
for j in range(_M[1]):
sid += 1
parent = create_linked_ball(sid, (dx*i-offset, dx*j-offset, 1e-3), '0 0 0', (2e-3, 0.9))
#child = create_linked_ball(sid + _M[0]*_M[1], (0, 0, 2e-3), '0 0 0', (1e-3, 0.9))
#parent.append(child)
sensor_pad.append(parent)
# Define tendons
tid = 0
for m in range(_M[0]):
for n in range(_M[1]):
tid += 1
if n != _M[0] - 1:
site_ids = [n + 1 + k + m*_M[0] for k in range(2)]
spatial = create_tendon(tid, ['tendon_site' + str(i) for i in site_ids])
tendon.append(spatial)
if m != _M[0] - 1:
tid += 1
site_ids = [n + 1 + (m + k)*_M[0] for k in range(2)]
spatial = create_tendon(tid, ['tendon_site' + str(i) for i in site_ids])
tendon.append(spatial)
tree.write(robot_xml) # Save robot
# Modify the keyframe home key qpos
njoints = np.prod(_M) + 10
scene_xml = 'models/scene.xml'
tree = et.parse(scene_xml)
root = tree.getroot()
key = root.find('.//keyframe/key')
qpos_value = '0 0 0 -1.570796 0 1.570796 0.785398 0 0 0'
key.attrib['ctrl'] = qpos_value
if njoints > 10:
qpos_value += ' ' + ' '.join(['0' for i in range(njoints - 10)])
key.attrib['qpos'] = qpos_value
tree.write(scene_xml)
# Load scene
scene_xml = 'models/scene.xml'
physics = mujoco.Physics.from_xml_path(scene_xml)
# Define simulation variables
site_name = 'attachment_site'
control_site = physics.data.site(name=site_name)
reach_sites = ['reach_site' + str(i+1) for i in range(3)]
joint_names = ['joint{}'.format(i+1) for i in range(7)]
def wrap2pi(xs):
return np.mod(xs, 2*np.pi)
def reach_test(target_name, duration=5.0):
omega = np.array([1, 1, 1])*2*np.pi/duration # Rotator angular velocities
target_site = physics.data.site(name=target_name)
video = []
stream = []
# Simulate, saving video frames
physics.reset(0)
physics.step()
ctrl = np.zeros(10)
quat = np.zeros(4)
mjlib.mju_mat2Quat(quat, control_site.xmat)
gain = 1.0/duration
with alive_bar(int(duration/_TIME_STEP), force_tty=True, title='Rendering') as bar:
while physics.data.time < duration:
# Compute the inverse kinematics
result = ik.qpos_from_site_pose(physics,
site_name,
target_pos=target_site.xpos,
target_quat=quat,
joint_names=joint_names,
tol=_TOL,
max_steps=_MAX_STEPS,
inplace=False
)
ctrl[:7] = result.qpos[:7]
ctrl[-3:] = wrap2pi(omega*physics.data.time)
physics.set_control(ctrl)
physics.step()
# Save contact pressure > 0
if physics.data.time >= 2.0:
pressure = physics.data.qpos[7:-3].copy().reshape(_M)
stream.append(np.expand_dims(pressure, -1))
# Save video frames and sensor data
if len(video) < physics.data.time * framerate:
pixels = physics.render(camera_id='prospective', width=width, height=height)
video.append(pixels.copy())
# Update progress
bar()
return video, stream
video1, stream1 = reach_test(reach_sites[0])
display_video(video1)
Rendering |████████████████████████████████████████| 2500/2500 [100%] in 16:50.0 (2.48/s)
display_video(stream1, 30, (5, 5))
video2, stream2 = reach_test(reach_sites[1])
display_video(video2)
Rendering |████████████████████████████████████████| 2500/2500 [100%] in 16:16.9 (2.56/s)
display_video(stream2, 30, (5, 5))
video3, stream3 = reach_test(reach_sites[2])
display_video(video3)
Rendering |████████████████████████████████████████| 2500/2500 [100%] in 18:46.4 (2.22/s)
display_video(stream3, 30, (5, 5))